# Code for dense filter

import jax
import jax.numpy as jnp
import numpy as np
import haiku as hk

from typing import Optional, Type, Union, List, Tuple
from jax.numpy import ndarray
from jax.nn.initializers import xavier_uniform, zeros
from .filter_base import Filter, _filter_register
from cybertron.utils.base import xaiver_uniform_init

@_filter_register("dense")
class DenseFilter(Filter):
    r"""Dense filter network.

    Args:
        dim_in (int):    Number of basis functions.
        dim_out (int):   Dimension of output filter Tensor.
        activation (hk.Module or function): Activation function. Default: None.
        n_hidden (int):  Number of hidden layers. Default: 1.
        name (str):      Name of the filter network. Default: "dense_filter".

    """

    def __init__(self,
                 dim_in: int,
                 dim_out: int,
                 activation = jax.nn.relu,
                 n_hidden: int = 1,
                 name: str = "dense_filter",
                 ):
        super().__init__(dim_in=dim_in,
                         dim_out=dim_out,
                         activation=activation,
                         name=name,
                         )
        
        self.n_hidden = n_hidden
        if self.n_hidden > 0:
            hidden_layers = [self.dim_out for _ in range(self.n_hidden + 1)]
            self.layers = hk.nets.MLP(output_sizes=hidden_layers,
                                   activation=self.activation,
                                   activate_final=False,
                                   w_init=xaiver_uniform_init,
                                   name="mlp",) # type: ignore
        else:
            self.layers = hk.nets.MLP(output_sizes=[self.dim_out],
                                      activation=self.activation,
                                      activate_final=True,
                                      w_init=xaiver_uniform_init,
                                      name="mlp",) # type: ignore
        
    def __call__(self, x: ndarray) -> ndarray:
        r"""Compute filter.
        Args:
            x (ndarray):    Array of shape (...,). Input array.
        """
        return self.layers(x)